gfx942: add flat + 16x FMoE kernels for GLM5.2-FP8 decode - #4040

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4040 --add-label <label>

Tested with AITER_LOG_MORE=1 python3 -u op_tests/test_moe_2stage.py -t 4 -d bf16 -dim 6144,256 -q 5 -a silu -e 257 -k 9 -s f --no-flydsl-csv.

  t=4
  ┌────────┬────────────────┬──────────────┐
  │  Run   │ ASM (flat_pf3) │ CK (2-stage) │
  ├────────┼────────────────┼──────────────┤
  │ 1      │ 59.0 us        │ 64.4 us      │
  ├────────┼────────────────┼──────────────┤
  │ 2      │ 56.5 us        │ 67.0 us      │
  ├────────┼────────────────┼──────────────┤
  │ 3      │ 57.7 us        │ 64.6 us      │
  ├────────┼────────────────┼──────────────┤
  │ 4      │ 56.9 us        │ 67.6 us      │
  ├────────┼────────────────┼──────────────┤
  │ 5      │ 55.1 us        │ 66.5 us      │
  ├────────┼────────────────┼──────────────┤
  │ median │ 56.9 us        │ 66.5 us      │
  └────────┴────────────────┴──────────────┘
  t=8
  ┌────────┬────────────────┬──────────────┐
  │  Run   │ ASM (flat_pf3) │ CK (2-stage) │
  ├────────┼────────────────┼──────────────┤
  │ 1      │ 78.0 us        │ 100.0 us     │
  ├────────┼────────────────┼──────────────┤
  │ 2      │ 82.1 us        │ 97.2 us      │
  ├────────┼────────────────┼──────────────┤
  │ 3      │ 80.7 us        │ 99.8 us      │
  ├────────┼────────────────┼──────────────┤
  │ 4      │ 112.1 us       │ 98.2 us      │
  ├────────┼────────────────┼──────────────┤
  │ 5      │ 78.6 us        │ 102.0 us     │
  ├────────┼────────────────┼──────────────┤
  │ median │ 80.7 us        │ 99.8 us      │
  └────────┴────────────────┴──────────────┘  

Conclusions:

• ASM is 1.17x faster at t=4.
• ASM is 1.24x faster at t=8 (ignoring the run 4 outlier at 112 us).
• precision: logits_diff is very small for all ASM runs

Performance results vs main (post PR #3422)

After #3422 ("GLM-5.1 tunings") merged, the CK 2-stage kernels for this workload were replaced by ASM 1-stage kernels on main. We repeated the comparison against the new baseline.

Note: "1-stage" means the full MOE pipeline — GEMM0 (gate) + GEMM1 (up) + SiLU + GEMM2 (down) — runs in a single kernel launch, as opposed to the old CK "2-stage" which used separate launches for stage1 and stage2.

Kernels used

Full MOE pipeline (profiler avg us/iter, gfx942 304 CUs)

Breakdown of gains

Accuracy (logits_diff vs BF16 reference)

Workload: GLM-5, E=257, topk=9, D=6144, D_FF=256, blockscale FP8, gfx942.

Updated performance results (50 profiler iterations, clean GPU)

Previous measurements used only 5 profiler iterations, which caused noisy/bimodal results. Re-ran both ASM flat and main with 50 profiler iterations on a clean GPU (no other workloads), 5 independent runs each.

Workload: GLM-5 FMoE decode, E=257, topk=9, dim=6144x256, A8W8 blockscale FP8, SiLU, gfx942

Device time (profiler, ROCTracer HW timestamps)

cuda.Event time (end-to-end GPU including host dispatch overhead)

Summary

Both main and ASM flat use 1-stage kernels (full MOE pipeline: GEMM0 gate + GEMM1 up + SiLU + GEMM2 down in a single kernel launch). The ASM flat kernel uses FLAT instructions for global memory access and reads activations in BF16 (X_BF16) instead of FP8, with a 16x128 tile size vs main's 32x128.